from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableParallel, RunnableBranch


def classify_question(input_dict):
    question = input_dict["question"]
    if (word in question for word in ["技术", "博客", "代码", "算法"]):
        return "tech"
    elif (word in question for word in ["商业", "市场", "投资", "利润"]):
        return "business"
    else:
        return "general"


# PromptTemplate.from_template返回的是一个<class 'langchain_core.prompt_values.StringPromptValue'>，可同text拿到文本
# prompt_template_parser = lambda x: print(type(x), x, x.text, "\n-----------------")
prompt_template_parser = lambda x: x.text

tech_chain = (
    PromptTemplate.from_template("你是一位技术专家，请用专业的技术知识回答：{question}")
    | prompt_template_parser
)

tech_chain.invoke({"question": "算法学习的建议"})

business_chain = (
    PromptTemplate.from_template("你是一位投资专家，请用专业的技术知识回答：{question}")
    | prompt_template_parser
)

general_chain = (
    PromptTemplate.from_template(
        "你是一位知识丰富的智能助手，请用专业的技术知识回答：{question}"
    )
    | prompt_template_parser
)

# RunnableBranch必须有defalut
select_runnable = RunnableBranch(
    (lambda x: classify_question(x) == "tech", tech_chain),
    (lambda x: classify_question(x) == "business", business_chain),
    general_chain,
)

# RunnableParallel返回一个字典
last_runnable = RunnableParallel(
    content=lambda x: x["question"], select_result=select_runnable
)

temp = last_runnable.invoke({"question": "算法学习的建议"})
print(type(temp), temp)
